"""数据库封装"""
import asyncio
from typing import Dict, Any

from contextlib import asynccontextmanager


from sqlalchemy.ext.asyncio import (
    create_async_engine, AsyncSession, async_sessionmaker, async_scoped_session
)
# from sqlalchemy.orm import DeclarativeBase
from sqlalchemy.orm import class_mapper

# 数据库配置
DATABASE_URL = "mysql+aiomysql://root:Da112871@47.113.101.61:3306/mp_util?charset=utf8"

# 创建异步引擎
async_engine = create_async_engine(
    DATABASE_URL,
    pool_size=20,  # 适当增加连接池大小
    max_overflow=40,  # 调整最大溢出连接数
    pool_recycle=3600,  # 调整连接回收时间
    pool_pre_ping=True  # 在使用连接之前先检查连接是否可用
)

# 创建异步会话工厂
async_session = async_sessionmaker(
    async_engine,
    class_=AsyncSession,
    expire_on_commit=False
)


@asynccontextmanager
async def session_context():
    """会话获取"""
    session = async_scoped_session(
        async_session,
        scopefunc=asyncio.current_task
    )
    try:
        yield session
        await session.commit()
    except:
        await session.rollback()
        raise
    finally:
        await session.close()
        await session.remove()


class BaseEx:
    """model 自定义属性扩展"""
    def to_dict(self) -> Dict[str, Any]:
        """将模型对象转换为字典"""
        # 获取模型对象的列名
        columns = [c.key for c in class_mapper(self.__class__).columns]
        # 将模型对象的列名和对应的值转换为字典
        return {c: getattr(self, c) for c in columns}
